#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <unistd.h>
#include <string.h>
#include <fcntl.h>
#include <getopt.h>
#include <errno.h>

#include "mbx_image.h"

#define TOOL_VERSION	"V1.3"

#define xprintf(fmt, ...)	\
	do { \
		if(bVerbose) printf(fmt, ##__VA_ARGS__); \
	} while(0)

#define ALIGIN(sz, a) \
	(((sz) + (a) - 1) & (~((a) - 1)))

#define ALIGIN_8(sz)	ALIGIN(sz, 8)
#define ALIGIN_16(sz)	ALIGIN(sz, 16)
#define ALIGIN_32(sz)	ALIGIN(sz, 32)

static bool bVerbose = false;

static const char soptions[] = "hvVS:m:u:x:L:A:B:r:R:a:d:c:p:i:o:l:";
static const struct option loptions[] = {
	{"help", 0, 0, 'h'},
	{"version", 0, 0, 'V'},
	{"verbose", 0, 0, 'v'},
	{"verify", 1, 0, 'S'},
	{"mbr", 1, 0, 'm'},
	{"uboot", 1, 0, 'u'},
	{"xlogo", 1, 0, 'x'},
	{"lboot", 1, 0, 'L'},
	{"aboot", 1, 0, 'A'},
	{"rboot", 1, 0, 'B'},
	{"lroot", 1, 0, 'r'},
	{"aroot", 1, 0, 'R'},
	{"lapps", 1, 0, 'a'},
	{"adata", 1, 0, 'd'},
	{"cache", 1, 0, 'c'},
	{"devic", 1, 0, 'p'},
	{"idisk", 1, 0, 'i'},
	{"list", 1, 0, 'l'},
	{"output", 1, 0, 'o'},
	{0, 0, 0, 0}
};

static const char *enames[] = {
	"MBR",
	"UBOOT", "XLOGO", "LBOOT", "ABOOT",
	"RBOOT", "LROOT", "AROOT", "LAPPS",
	"ADATA", "CACHE", "DEVIC", "IDISK",
	"Unknow"
};

static void *load_file(const char *fn, unsigned *_sz)
{
    char *data;
    int sz;
    int fd;

	if(!fn) return 0;

	xprintf(" == Loading file : %s  [", fn);

    data = 0;
    fd = open(fn, O_RDONLY);
    if(fd < 0) {
		xprintf("ERROR]\n");
		printf("Ignore Invalid input file : %s\n", fn);
		return 0;
	}

    sz = lseek(fd, 0, SEEK_END);
    if(sz < 0) goto oops;

    if(lseek(fd, 0, SEEK_SET) != 0) goto oops;

    data = (char*) malloc(sz);
    if(data == 0) goto oops;

    if(read(fd, data, sz) != sz) goto oops;
    close(fd);

    if(_sz) *_sz = sz;
	xprintf("OK]\n");
    return data;

oops:
	xprintf("ERROR]\n");
	printf("Ignore Invalid input file : %s\n", fn);
    close(fd);
    if(data != 0) free(data);
    return 0;
}

void show_usage(char *arg0)
{
	printf("Usage : %s options\n", arg0);
	printf("  -h,--help            show this message.\n");
	printf("  -V,--version         show tool version.\n");
	printf("  -v,--verbose         show more information.\n");
	printf("  -S,--verify          set verify for board specific value.\n");
	printf("  -m,--mbr 'file'      pack 'file' as mbr partition table.\n");
	printf("  -u,--uboot 'file'    pack 'file' as bootloader.\n");
	printf("  -x,--xlogo 'file'    pack 'file' as logo.\n");
	printf("  -L,--lboot 'file'    pack 'file' as linux kernel.\n");
	printf("  -A,--aboot 'file'    pack 'file' as anrdoid boot.\n");
	printf("  -B,--rboot 'file'    pack 'file' as anrdoid recovery.\n");
	printf("  -r,--lroot 'file'    pack 'file' as linux rootfs.\n");
	printf("  -R,--aroot 'file'    pack 'file' as anrdoid system.\n");
	printf("  -a,--lapps 'file'    pack 'file' as linux apps.\n");
	printf("  -d,--adata 'file'    pack 'file' as android data.\n");
	printf("  -c,--cache 'file'    pack 'file' as android cache.\n");
	printf("  -p,--devic 'file'    pack 'file' as android device.\n");
	printf("  -i,--idisk 'file'    pack 'file' as sdcard.\n");
	printf("  -l,--list 'file'     list mbx image information.\n");
	printf("  -o,--output 'file'   create 'file' as target image.\n");
}

void dump_hex16(const char *prefix, unsigned char *hex16)
{
	int i;
	printf("%s", prefix);
	for(i = 0; i < 16; ++i) {
		printf("%02X ", hex16[i]);
		// if(i == 7) printf(" ");
	}

	printf("| ");
	for(i = 0; i < 16; ++i) {
		if(hex16[i] < 0x20) printf(".");
		else if(hex16[i] > 0x7E) printf(".");
		else printf("%c", hex16[i]);
	}
	printf("\n");
}

void list_info(char *filename)
{
	int fd, i, sz, res;
	unsigned char headdump[16];
	struct mbx_header hdr;
	struct entry_info* pentry = NULL;
	fd = open(filename, O_RDONLY);
	if(fd < 0) {
		perror("open");
		return;
	}

	sz = lseek(fd, 0, SEEK_END);
    lseek(fd, 0, SEEK_SET);

	res = read(fd, &hdr, sizeof(hdr));
	if(res != sizeof(hdr)) {
		perror("read");
		goto done;
	}
	
	if(hdr.magic != MBX_MAGIC) {
		fprintf(stderr, "Not a MBX image file.\n");
		goto done;
	}
	
	if(sz != hdr.mbx_length) {
		fprintf(stderr, "Except file eof.\n");
		goto done;
	}

	printf("mbx image dummp [%s]:\n", filename);
	printf("  len     : %u\n", sz);
	printf("  version : %u.%u\n", MBX_VER_MAJOR(hdr.version),
		MBX_VER_MINOR(hdr.version));
	printf("  verify  : 0x%08x\n", hdr.verify);
	printf("  entrys  : %u\n", hdr.entrys);
	
	sz = sizeof(struct entry_info) * hdr.entrys;
	
	pentry = (struct entry_info*)malloc(sz);
	read(fd, pentry, sz);
	for(i = 0; i < hdr.entrys; ++i) {
		printf("   [%d] %s : offset=0x%08X, size=%u, head dump:\n", i, pentry[i].name,
			pentry[i].offset, pentry[i].length);
		lseek(fd, pentry[i].offset, SEEK_SET);
		read(fd, headdump, 16);
		dump_hex16("       ", headdump);
		read(fd, headdump, 16);
		dump_hex16("       ", headdump);
	}
	free(pentry);
done:
	close(fd);
}

int main(int argc, char *argv[])
{
	char c;
	int opt_idx;
	int i, cnt, res, fd, ensz, bso;
	unsigned int verify = 0;
	struct mbx_header hdr;
	struct entry_info entry[MBX_ENTRY_TYPE_MAX];	// idx = 3 is not used
	void *pbuf[MBX_ENTRY_TYPE_MAX] = { 0 };
	char *pfn[MBX_ENTRY_TYPE_MAX] = { 0 };
	unsigned int bufsz[MBX_ENTRY_TYPE_MAX] = { 0 };
	char padding[32];
	bzero(padding, 32);

	char *outmbx_fn = NULL;

	while ((c = getopt_long(argc, argv, soptions, loptions, &opt_idx)) != -1) {
		switch (c) {
		case 'h':
			show_usage(argv[0]);
			return 0;
		case 'V':
			printf("version : %s\n", TOOL_VERSION);
			return 0;
		case 'l':
			list_info(optarg);
			return 0;
		case 'S':
			verify = atoi(optarg);
			break;
		case 'v': bVerbose = true; break;
		case 'o': outmbx_fn = optarg; break;
		case 'm': pfn[MBX_ENTRY_TYPE_MBR] = optarg; break;
		case 'u': pfn[MBX_ENTRY_TYPE_UBOOT] = optarg; break;
		case 'x': pfn[MBX_ENTRY_TYPE_XLOGO] = optarg; break;
		case 'L': pfn[MBX_ENTRY_TYPE_LBOOT] = optarg; break;
		case 'A': pfn[MBX_ENTRY_TYPE_ABOOT] = optarg; break;
		case 'B': pfn[MBX_ENTRY_TYPE_RBOOT] = optarg; break;
		case 'r': pfn[MBX_ENTRY_TYPE_LROOT] = optarg; break;
		case 'R': pfn[MBX_ENTRY_TYPE_AROOT] = optarg; break;
		case 'a': pfn[MBX_ENTRY_TYPE_LAPPS] = optarg; break;
		case 'd': pfn[MBX_ENTRY_TYPE_ADATA] = optarg; break;
		case 'c': pfn[MBX_ENTRY_TYPE_CACHE] = optarg; break;
		case 'p': pfn[MBX_ENTRY_TYPE_DEVIC] = optarg; break;
		case 'i': pfn[MBX_ENTRY_TYPE_IDISK] = optarg; break;
		default:
			fprintf(stderr, "Try `%s --help' for more information.\n", argv[0]);
			return 1;
		}
	}

	if(!outmbx_fn) {
		show_usage(argv[0]);
		return 1;
	}

	bzero(entry, sizeof(entry));
	bzero(bufsz, sizeof(bufsz));
	bzero(&hdr, sizeof(hdr));
	hdr.magic = MBX_MAGIC;
	hdr.version = MBX_MAKE_VERSION(1, 0);
	hdr.mbx_length = sizeof(hdr);
	memcpy(&hdr.verify, &verify, 4);

	cnt = 0;
	ensz = sizeof(entry[0]);
	for(i = 0; i < MBX_ENTRY_TYPE_MAX; ++i) {
		if(!pfn[i]) continue;
		pbuf[i] = load_file(pfn[i], &bufsz[i]);
		if(bufsz[i] > 0) {
			strcpy(entry[i].name, enames[i]);
			entry[i].type = i;
			entry[i].length = bufsz[i];
			hdr.mbx_length += ALIGIN_32(bufsz[i]);
			hdr.mbx_length += ensz;
			cnt++;
		}
	}

	bso = sizeof(hdr) + ensz * cnt;
	for(i = 0; i < MBX_ENTRY_TYPE_MAX; ++i) {
		if(!bufsz[i]) continue;
		entry[i].offset = bso;
		bso += ALIGIN_32(entry[i].length);
	}

	assert(hdr.mbx_length == bso);

	hdr.entrys = cnt;
	xprintf("Verify ID : %u.\n", hdr.verify);
	xprintf("Will pack %d images.\n", cnt);
	fd = open(outmbx_fn, O_RDWR | O_CREAT | O_TRUNC, 0664);
	res = write(fd, &hdr, sizeof(hdr));
	if(res != sizeof(hdr)) goto err;

	for(i = 0; i < MBX_ENTRY_TYPE_MAX; ++i) {
		if(!bufsz[i]) continue;
		write(fd, &entry[i], sizeof(entry[i]));
	}

	for(i = 0; i < MBX_ENTRY_TYPE_MAX; ++i) {
		if(!bufsz[i]) continue;
		xprintf(" == '%s' sz=%d : %s\n", enames[i], bufsz[i], pfn[i]);
		res = write(fd, pbuf[i], bufsz[i]);
		if(res != bufsz[i]) goto err;

		int psz = ALIGIN_32(bufsz[i]) - bufsz[i];
		if(psz > 0) {
			write(fd, padding, psz);
		}
		free(pbuf[i]);
	}

	printf("Create mbx image file '%s' done!\n", outmbx_fn);

	close(fd);
	return 0;
err:
	fprintf(stderr, "Write data error : %s\n", strerror(errno));
	close(fd);
	unlink(outmbx_fn);
	return 1;
}
